use std::any::Any;
use std::future::Future;
use std::sync::Arc;

use hashbrown::HashMap;
use tokio::sync::Mutex;

// 单飞模式, 模拟 go 的 singleflight 写的
// 用于避免重复请求, 当有相同的请求时, 直接返回之前的结果
// 适用于耗时长的请求, 如网络请求, 数据库查询等
// 需要拥有多线程的内部可变性, 所以需要 Arc + Mutex
// 需要支持不同的类型分发到不同的线程, 所以要用 Arc<dyn Any + Send + Sync>
// type Call<T: Any + Send + Sync> = Arc<Mutex<Arc<T>>>;
type ResultOptionArcAny = Result<Option<Arc<dyn Any + Send + Sync>>, Arc<dyn Any + Send + Sync>>;

#[derive(Debug, Default, Clone)] // Group 是需要给到不同的线程的, 所以需要 Clone
pub struct AsyncGroup {
    // 这里加个 Arc + Mutex 包装一下, 使得内部可变性可以被多个线程访问
    m: Arc<Mutex<HashMap<String, Arc<Mutex<ResultOptionArcAny>>>>>,
}

impl AsyncGroup {
    pub async fn work<T, F, Fut, E>(&self, key: &str, func: F) -> Result<Option<Arc<T>>, Arc<E>>
        where
            T: Any + Send + Sync,
            E: Any + Send + Sync,
            Fut: Future<Output=Result<Option<T>, E>>,
            F: FnOnce() -> Fut,
    {
        let mut m = self.m.lock().await;
        if let Some(c) = m.get(key) {
            let c = c.clone();
            drop(m);
            let res = c.lock().await.clone();
            return res
                .map(|o| o.map(|v| v.downcast::<T>().unwrap()))
                .map_err(|e| e.downcast::<E>().unwrap()); // 这里类型转换肯定是Ok的
        }

        let call = Arc::new(Mutex::new(Ok(None)));
        let call_clone = call.clone();
        m.insert(key.to_owned(), call_clone);
        let mut call_guard = call.lock().await;
        drop(m);
        let res = match func().await {
            Ok(v) => match v {
                Some(x) => Ok(Some(Arc::new(x))),
                None => Ok(None),
            },
            Err(e) => Err(Arc::new(e)),
        };
        let res_clone = res.clone();
        *call_guard = res_clone
            .map(|x| x.map(|v| v as _))
            .map_err(|x| x as _);
        drop(call_guard);
        self.m.lock().await.remove(key).unwrap();

        res
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::time::Duration;

    use crate::async_single_flight::AsyncGroup;

    #[derive(Clone, Debug)]
    struct A {
        value: i32,
    }

    #[test]
    fn test_single_flight() {
        let body = async {
            let count = 10;
            let (tx, rx) = mpsc::channel::<()>();
            let g = AsyncGroup::default();
            for i in 0..count {
                let g = g.clone();
                let tx = tx.clone();
                tokio::spawn(async move {
                    let (key, value) = if i % 2 == 0 {
                        ("aaa", 121)
                    } else {
                        ("bbb", 111)
                    };
                    let res = g
                        .work(key, || async move {
                            tokio::time::sleep(Duration::from_millis(300)).await;
                            Ok::<Option<i32>, String>(Some(value))
                        })
                        .await;
                    let res_value = res.unwrap().unwrap();
                    assert_eq!(*res_value, value);
                    drop(tx);
                });
            }

            for i in 0..count {
                let g = g.clone();
                let tx = tx.clone();
                tokio::spawn(async move {
                    let (key, value) = if i % 2 == 0 {
                        ("ccc", "aaa".to_string())
                    } else {
                        ("ddd", "bbb".to_string())
                    };
                    let value_clone = value.clone();
                    let res = g
                        .work(key, || async move {
                            tokio::time::sleep(Duration::from_millis(300)).await;
                            Ok::<Option<String>, String>(Some(value_clone))
                        })
                        .await;
                    let res_value = res.unwrap().unwrap();
                    assert_eq!(*res_value, value);
                    drop(tx);
                });
            }

            for i in 0..count {
                let g = g.clone();
                let tx = tx.clone();
                tokio::spawn(async move {
                    let (key, value) = if i % 2 == 0 {
                        ("eee", A { value: 121 })
                    } else {
                        ("fff", A { value: 111 })
                    };
                    let value_clone = value.clone();
                    let res = g
                        .work(key, || async move {
                            tokio::time::sleep(Duration::from_millis(300)).await;
                            Ok::<Option<A>, String>(Some(value_clone))
                        })
                        .await;
                    let res_value = res.unwrap().unwrap();
                    assert_eq!(res_value.value, value.value);
                    drop(tx);
                });
            }

            drop(tx);
            while rx.recv().is_ok() {}
        };

        tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .unwrap()
            .block_on(body);
    }
}
